import cupy as cp
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sg
import random


# 作图设置
plt.rcParams['axes.unicode_minus'] = False
plt.xticks(fontproperties='Times New Roman')


class ASR:
    # 通用参数
    Re = 6371e3  # 地球半径
    c = 3e8  # 光速
    G = 6.67e-11  # 引力常量
    M = 5.965e24  # 地球质量

    def __init__(self, antenna_parameter: dict, radar_parameter: dict, mission_parameter: dict, sample_parameter: dict):
        # 天线参数读入
        self.Nx = antenna_parameter['Nx']
        self.Ny = antenna_parameter['Ny']
        self.Wxt = antenna_parameter['Wxt']
        self.Wyt = antenna_parameter['Wyt']
        self.Wxr = antenna_parameter['Wxr']
        self.Wyr = antenna_parameter['Wyr']
        self.dxRatio = antenna_parameter['dxRatio']
        self.dyRatio = antenna_parameter['dyRatio']

        # 雷达参数读入
        self.H = radar_parameter['H']
        self.f0 = radar_parameter['f0']
        self.PRF = radar_parameter['PRF']
        self.Br = radar_parameter['Br']
        self.N = radar_parameter['N']

        # 任务参数读入
        self.alpha_c = mission_parameter['alphaC']
        self.theta_sq_c = mission_parameter['thetaSqC']
        self.alpha_scan_min = mission_parameter['alphaScanMin']
        self.alpha_scan_max = mission_parameter['alphaScanMax']
        self.alpha_c_r = np.deg2rad(self.alpha_c)
        self.theta_sq_c_r = np.deg2rad(self.theta_sq_c)
        self.alpha_scan_max_r = np.deg2rad(self.alpha_scan_max)
        self.alpha_scan_min_r = np.deg2rad(self.alpha_scan_min)

        # 采样手动设置
        self.R_sample_rate = sample_parameter['RSampleRate']
        self.f_sample_rate = sample_parameter['fSampleRate']

        # 雷达姿态
        self.A = np.arctan(np.tan(self.alpha_c_r) / np.cos(self.theta_sq_c_r))
        self.theta_yaw = np.arctan(np.sin(self.alpha_c_r) * np.tan(self.theta_sq_c_r))
        self.theta_pit = np.arcsin(np.cos(self.alpha_c_r) * np.sin(self.theta_sq_c_r))

        # 地球最大扫描角度
        self.theta_earth_r = np.arcsin(self.Re / (self.Re + self.H))
        self.theta_earth = np.rad2deg(self.theta_earth_r)

        # 雷达参数计算
        self.PRF_r = self.PRF * self.N
        self.pr = self.c / 2 / self.Br
        self.lamda = self.c / self.f0  # 波长
        self.v = (self.G * self.M / (self.H + self.Re)) ** 0.5  # 雷达平台速度

        # 天线参数
        self.dx = self.dxRatio * self.lamda  # 天线x单元间隔
        self.dy = self.dyRatio * self.lamda  # 天线y单元间隔

        # 多普勒计算
        self.f_eta_c = 2 * self.v * np.sin(self.theta_sq_c_r) / self.lamda  # 多普勒中心频率
        self.theta_bw_r = 0.886 / self.dyRatio / self.Ny / 4  # 波束宽度
        self.theta_bw = np.rad2deg(self.theta_bw_r)
        self.Bd = 2 * self.v * np.cos(self.theta_sq_c_r) * self.theta_bw_r / self.lamda
        # self.Bd = 2 * self.v * np.cos(self.theta_sq_c_r) / (self.Ny * self.dy)

        # 主信号区
        self.theta_min_r = self.cal_theta(self.alpha_scan_min_r, self.theta_sq_c_r)
        self.theta_max_r = self.cal_theta(self.alpha_scan_max_r, self.theta_sq_c_r)
        self.theta_min = np.rad2deg(self.theta_min_r)
        self.theta_max = np.rad2deg(self.theta_max_r)

        self.R_min = self.theta2r_eta_cpu(self.theta_min_r)[0]  # 最近斜距
        self.R_max = self.theta2r_eta_cpu(self.theta_max_r)[0]  # 最远斜距

        self.R0 = np.linspace(self.R_min, self.R_max, self.R_sample_rate)
        self.theta_nadir0_r = self.r2eta_theta_cpu(self.R0)[1]
        self.theta_nadir0 = np.rad2deg(self.theta_nadir0_r)

        # 模糊距离计算
        self.R_amb = []
        delta_r = self.c / 2 / self.PRF
        r_max = np.sqrt((self.Re + self.H) ** 2 - self.Re ** 2)

        # 测绘带到星下点（范围， 已包含镜像区）
        r = self.R0.copy()
        while True:
            r -= delta_r
            if r[0] < self.H and r[-1] < self.H:
                break
            elif r[0] < self.H < r[-1]:
                print('Warning: Nadir Echo')
                r[np.where(r < self.H)] = np.NAN
                self.R_amb.append(r.copy())
                break
            else:
                self.R_amb.append(r.copy())
        self.R_amb.reverse()

        # 测绘带到地球边界（已包含镜像区）
        r = self.R0.copy()
        while True:
            r += delta_r
            if r[-1] > r_max:
                break
            else:
                self.R_amb.append(r.copy())

    def doppler2theta_sq(self, fd):
        return cp.arcsin(self.lamda * fd / 2 / self.v)

    def r2eta_theta(self, r):  # 由斜距计算入射角，星下点角
        theta = cp.arccos(cp.round((r ** 2 + (self.Re + self.H) ** 2 - self.Re ** 2) / 2 / (self.Re + self.H) / r, 12))
        # alpha = np.arccos((r ** 2 + (Re + self.H) ** 2 - Re ** 2) / 2 / (Re + self.H) / r)
        eta = cp.arcsin(np.round(cp.sin(theta) * (self.Re + self.H) / self.Re, 12))
        # eta = np.arcsin(np.sin(alpha) * (Re + self.H) / Re)
        return eta, theta

    def r2eta_theta_cpu(self, r):  # 由斜距计算入射角，星下点角
        theta = np.arccos(np.round((r ** 2 + (self.Re + self.H) ** 2 - self.Re ** 2) / 2 / (self.Re + self.H) / r, 12))
        # alpha = np.arccos((r ** 2 + (Re + self.H) ** 2 - Re ** 2) / 2 / (Re + self.H) / r)
        eta = np.arcsin(np.round(np.sin(theta) * (self.Re + self.H) / self.Re, 12))
        # eta = np.arcsin(np.sin(alpha) * (Re + self.H) / Re)
        return eta, theta

    def theta2r_eta(self, theta):  # 由星下点角计算斜距、入射角
        eta = cp.arcsin(cp.round(np.sin(theta) * (self.Re + self.H) / self.Re, 12))
        # eta = np.arcsin(np.sin(alpha) * (Re + self.H) / Re)
        beta = eta - theta
        r = cp.sin(beta) / cp.sin(theta) * self.Re
        if type(r) == cp.float64 and r == 0:
            r = self.H
        elif type(r) == cp.array:
            r[np.where(beta == 0)] = self.H
        return r, eta

    def theta2r_eta_cpu(self, theta):  # 由星下点角计算斜距、入射角
        eta = np.arcsin(np.round(np.sin(theta) * (self.Re + self.H) / self.Re, 12))
        # eta = np.arcsin(np.sin(alpha) * (Re + self.H) / Re)
        beta = eta - theta
        r = np.sin(beta) / np.sin(theta) * self.Re
        if type(r) == np.float64 and r == 0:
            r = self.H
        elif type(r) == np.array:
            r[np.where(beta == 0)] = self.H
        return r, eta

    def steer(self, alpha, theta_sq):
        def phase(al, th):
            phx = ((cp.cos(self.A) * cp.cos(self.theta_yaw) +
                    cp.sin(self.A) * cp.sin(self.theta_yaw) * cp.sin(self.theta_pit)) * cp.sin(al) * cp.cos(th) +
                   (cp.cos(self.A) * cp.sin(self.theta_yaw) -
                    cp.sin(self.A) * cp.cos(self.theta_yaw) * cp.sin(self.theta_pit)) * cp.sin(th) -
                   cp.sin(self.A) * cp.cos(self.theta_pit) * cp.cos(al) * cp.cos(th)) * \
                  self.dx * 2j * cp.pi / self.lamda

            phy = (-cp.cos(self.theta_pit) * cp.sin(self.theta_yaw) * cp.sin(al) * cp.cos(th) +
                   cp.cos(self.theta_pit) * cp.cos(self.theta_yaw) * cp.sin(th) -
                   cp.sin(self.theta_pit) * cp.cos(al) * cp.cos(th)) * self.dy * 2j * cp.pi / self.lamda

            phz = ((cp.sin(self.A) * cp.cos(self.theta_yaw) -
                    cp.cos(self.A) * cp.sin(self.theta_yaw) * cp.sin(self.theta_pit)) * cp.sin(al) * cp.cos(th) +
                   (cp.sin(self.A) * cp.sin(self.theta_yaw) +
                    cp.cos(self.A) * cp.cos(self.theta_yaw) * cp.sin(self.theta_pit)) * cp.sin(th) +
                   cp.cos(self.A) * cp.cos(self.theta_pit) * cp.cos(al) * cp.cos(th))
            return phx, phy, phz

        alpha_m, theta_sq_m = cp.meshgrid(alpha, theta_sq)
        phase_x, phase_y, phase_z = phase(alpha_m, theta_sq_m)

        shape_x = list(phase_x.shape)
        shape_y = list(phase_y.shape)
        shape_p = list(phase_x.shape)
        shape_x.extend([self.Nx, 1])
        shape_y.extend([self.Ny, 1])

        sx = cp.zeros(tuple(shape_x), dtype='complex')
        sy = cp.zeros(tuple(shape_y), dtype='complex')
        p = cp.zeros(tuple(shape_p), dtype='complex')
        for i in range(4):
            p += cp.exp(i * phase_y)
        for i in range(self.Nx):
            sx[..., i, 0] = cp.exp(i * phase_x)
        for i in range(self.Ny):
            sy[..., i, 0] = p * cp.exp(4 * i * phase_y)

        loc = np.where(phase_z < 0)
        sx[loc[0], loc[1], ...] = cp.nan + cp.nan * 1j
        sy[loc[0], loc[1], ...] = cp.nan + cp.nan * 1j

        return sx, sy

    def steer_2d(self, alpha, theta_sq):
        def phase(al, th):
            phx = ((cp.cos(self.A) * cp.cos(self.theta_yaw) +
                    cp.sin(self.A) * cp.sin(self.theta_yaw) * cp.sin(self.theta_pit)) * cp.sin(al) * cp.cos(th) +
                   (cp.cos(self.A) * cp.sin(self.theta_yaw) -
                    cp.sin(self.A) * cp.cos(self.theta_yaw) * cp.sin(self.theta_pit)) * cp.sin(th) -
                   cp.sin(self.A) * cp.cos(self.theta_pit) * cp.cos(al) * cp.cos(th)) * \
                  self.dx * 2j * cp.pi / self.lamda

            phy = (-cp.cos(self.theta_pit) * cp.sin(self.theta_yaw) * cp.sin(al) * cp.cos(th) +
                   cp.cos(self.theta_pit) * cp.cos(self.theta_yaw) * cp.sin(th) -
                   cp.sin(self.theta_pit) * cp.cos(al) * cp.cos(th)) * self.dy * 2j * cp.pi / self.lamda

            phz = ((cp.sin(self.A) * cp.cos(self.theta_yaw) -
                    cp.cos(self.A) * cp.sin(self.theta_yaw) * cp.sin(self.theta_pit)) * cp.sin(al) * cp.cos(th) +
                   (cp.sin(self.A) * cp.sin(self.theta_yaw) +
                    cp.cos(self.A) * cp.cos(self.theta_yaw) * cp.sin(self.theta_pit)) * cp.sin(th) +
                   cp.cos(self.A) * cp.cos(self.theta_pit) * cp.cos(al) * cp.cos(th))
            return phx, phy, phz

        phase_x, phase_y, phase_z = phase(alpha, theta_sq)

        shape_x = list(phase_x.shape)
        shape_y = list(phase_y.shape)
        shape_x.extend([self.Nx, 1])
        shape_y.extend([self.Ny, 1])
        sx = cp.zeros(tuple(shape_x), dtype='complex')
        sy = cp.zeros(tuple(shape_y), dtype='complex')
        p = cp.zeros(phase_x.shape, dtype='complex')
        for i in range(4):
            p += cp.exp(i * phase_y)
        for i in range(self.Nx):
            sx[..., i, 0] = cp.exp(i * phase_x)
        for i in range(self.Ny):
            sy[..., i, 0] = p * cp.exp(4 * i * phase_y)

        loc = np.where(phase_z < 0)
        sx[loc[0], loc[1], ...] = cp.nan + cp.nan * 1j
        sy[loc[0], loc[1], ...] = cp.nan + cp.nan * 1j

        return sx, sy

    def vector_2d(self, alpha, theta_sq):
        sx, sy = self.steer_2d(alpha, theta_sq)
        v = self.vector(sx, sy)
        shape = list(alpha.shape)
        shape.extend([self.Nx * self.Ny, 1])
        v = v.reshape(shape)
        return v

    def e2gm(self, p, v):
        p = abs(p)
        p = p ** 2
        scale = 100
        start = 0
        end = start + scale
        m = cp.zeros((self.R_sample_rate, self.Ny, self.Nx * self.Ny), dtype='complex')
        # for i in range(v.shape[0]):
        #     m[i, ...] = cp.sum(p[i, ...] * (v[i, ...] @ np.conj(cp.transpose(v[i, ...], (0, 2, 1)))), axis=0)
        while end < self.R_sample_rate:
            m[start: end, ...] += cp.sum(p[start: end, ...] * (v[start: end, :, 0: self.Ny, :] @
                                                               np.conj(cp.transpose(v[start: end, ...], (0, 1, 3, 2)))), axis=1)
            start += scale
            end += scale
        m[start:, ...] += cp.sum(p[start:, ...] * (v[start:, :, 0: self.Ny, :] @
                                                   np.conj(cp.transpose(v[start:, ...], (0, 1, 3, 2)))), axis=1)
        return m

    def separated_pattern(self, alpha, theta_sq, wx, wy):  # w.shape = (N, 1) alpha.shape = (N, ) theta_sq.shape = (N, )
        sx, sy = self.steer(alpha, theta_sq)
        s = (cp.transpose(wx) @ sx) * (cp.transpose(wy) @ sy)
        return s.reshape(len(theta_sq), len(alpha))

    def plot_pattern(self, wx, wy):
        alpha = cp.linspace(self.alpha_c - 90, self.alpha_c + 90, 1801)
        # alpha = cp.linspace(-180, 180, 1801)
        theta_sq = cp.linspace(self.theta_sq_c - 90, self.theta_sq_c + 90, 1801)
        # theta_sq = cp.linspace(-180, 180, 1801)
        alpha_r, theta_sq_r = cp.deg2rad(alpha), cp.deg2rad(theta_sq)
        s = self.separated_pattern(alpha_r, theta_sq_r, wx, wy)
        s = abs(s)
        s = cp.asnumpy(s)
        s = s / np.nanmax(s)
        s_db = 20 * np.log10(s)
        s_db[np.where(s_db < -50)] = -50

        alpha, theta_sq = cp.asnumpy(alpha), cp.asnumpy(theta_sq)
        alpha, theta_sq = np.meshgrid(alpha, theta_sq)
        plt.contourf(alpha, theta_sq, s_db, cmap='jet')
        plt.colorbar()
        plt.xlabel(r'$\alpha$')
        plt.ylabel(r'$\theta_{sq}$')

    def cal_matrix(self):
        # 权重
        w = self.weight(self.Wxt, self.Wyt)
        w = w.reshape(-1, 1)

        # 多普勒频率
        f_eta_max = self.f_eta_c + self.Bd / 2
        f_eta_min = self.f_eta_c - self.Bd / 2
        f_eta0 = cp.linspace(f_eta_min, f_eta_max, self.f_sample_rate)

        # 模糊距离
        r0 = self.R0.copy()
        r_amb = self.R_amb.copy()
        r_amb.append(r0.copy())
        r0 = cp.asarray(r0)
        r_amb = np.array(r_amb)
        r_amb = cp.asarray(r_amb)

        # 形状及数目
        shape = [self.R_sample_rate, 1, 1]
        num = 0

        # 主信号区
        eta, theta = self.r2eta_theta(r0)
        theta_sq = self.doppler2theta_sq(f_eta0)
        alpha, theta_sq, theta = self.cal_alpha_all(theta_sq, theta)
        e = self.vector_2d(alpha, theta_sq)
        p = w.T @ e
        m = self.e2gm(p, e)
        sigma = self.cal_scatter(eta)
        r, eta, sigma = r0.reshape(shape), eta.reshape(shape), sigma.reshape(shape)
        s0 = sigma / r ** 3 / cp.sin(eta) * m
        print('Main Signal Area Finish')

        # 所有信号区
        s_amb = cp.zeros((self.R_sample_rate, self.Ny, self.Nx * self.Ny), dtype='complex')

        # 测绘带及测绘带向上
        f_eta = f_eta0.copy()
        while True:
            theta_sq = self.doppler2theta_sq(f_eta)
            for i in range(r_amb.shape[0]):
                # 同向
                eta, theta = self.r2eta_theta(r_amb[i, ...])
                alpha, theta_sq_m, theta = self.cal_alpha_all(theta_sq, theta)
                e = self.vector_2d(alpha, theta_sq_m)
                p = w.T @ e
                m = self.e2gm(p, e)
                sigma = self.cal_scatter(eta)
                r, eta, sigma = r_amb[i, ...].reshape(shape), eta.reshape(shape), sigma.reshape(shape)
                s_amb += cp.nan_to_num(sigma / r ** 3 / cp.sin(eta) * m)

                # 镜像
                alpha = -alpha
                e = self.vector_2d(alpha, theta_sq_m)
                p = w.T @ e
                m = self.e2gm(p, e)
                s_amb += cp.nan_to_num(sigma / r ** 3 / cp.sin(eta) * m)

                print(f'{num}:{i}', end='\r')

            # 停止条件
            num += 1
            f_eta += self.PRF_r
            if np.any(np.isnan(cp.asnumpy(theta_sq))):
                break

        # 测绘带向下
        f_eta = f_eta0.copy()
        while True:
            f_eta -= self.PRF_r
            theta_sq = self.doppler2theta_sq(f_eta)
            for i in range(r_amb.shape[0]):
                # 同向
                eta, theta = self.r2eta_theta(r_amb[i, ...])
                alpha, theta_sq_m, theta = self.cal_alpha_all(theta_sq, theta)
                e = self.vector_2d(alpha, theta_sq_m)
                p = w.T @ e
                m = self.e2gm(p, e)
                sigma = self.cal_scatter(eta)
                r, eta, sigma = r_amb[i, ...].reshape(shape), eta.reshape(shape), sigma.reshape(shape)
                s_amb += cp.nan_to_num(sigma / r ** 3 / cp.sin(eta) * m)

                # 镜像
                alpha = -alpha
                e = self.vector_2d(alpha, theta_sq_m)
                p = w.T @ e
                m = self.e2gm(p, e)
                s_amb += cp.nan_to_num(sigma / r ** 3 / cp.sin(eta) * m)

                print(f'{num}:{i}', end='\r')

            # 停止条件
            if np.any(np.isnan(cp.asnumpy(theta_sq))):
                break
            num += 1

        print(f'{num}:{i}')
        s_amb -= s0
        return s0, s_amb

    def e2m(self, s):
        m = np.zeros([s.shape[0], self.Nx * self.Ny, self.Nx * self.Ny], dtype='complex')
        t = s.copy()
        for i in range(self.Nx):
            m[:, self.Ny * i: self.Ny * (i + 1), :] = t
            if i < self.Nx - 1:
                t[:, :, self.Ny:] = t[:, :, 0: (self.Nx-1) * self.Ny]
                t[:, :, 0: self.Ny] = np.transpose(np.conj(s[:, :, self.Ny * (i + 1): self.Ny * (i + 2)]), (0, 2, 1))
        return m

    def plot_rar_r(self, ss, sa, wx, wy, if_plot=False):
        wmn = self.weight(wx, wy)
        wmn = wmn.reshape(self.Nx * self.Ny, 1)
        ss = self.e2m(ss)
        sa = self.e2m(sa)
        ss = wmn.T @ ss @ np.conj(wmn)
        sa = wmn.T @ sa @ np.conj(wmn)
        rar = (abs(sa) / abs(ss)).reshape(self.R_sample_rate, )
        rar_db = 10 * np.log10(rar)
        if if_plot:
            plt.plot(self.R0, rar_db)
            plt.xlabel('distance')
            plt.ylabel('RASR/dB')
        return rar_db

    @staticmethod
    def cal_theta(alpha, theta_sq):  # input:radian output:radian
        return np.arccos(np.cos(alpha) * np.cos(theta_sq))

    @staticmethod
    def cal_alpha_all(theta_sq, theta):
        theta_sq, theta = cp.meshgrid(theta_sq, theta)
        return cp.arccos(cp.cos(theta) / cp.cos(theta_sq)), theta_sq, theta

    @staticmethod
    # parameter from 《Handbook of Radar Scattering Statistics for Terrain>> Appendix C grass L band HV
    def cal_scatter(eta):  # input: radian output: power
        # HH
        p1 = -29.235
        p2 = 37.550
        p3 = 2.332
        p4 = -2.615
        p5 = 5.0
        p6 = -1.616

        # HV
        # p1 = -40.166
        # p2 = 26.833
        # p3 = 2.029
        # p4 = -1.473
        # p5 = 3.738
        # p6 = -1.324

        sigma0_db = p1 + p2 * cp.exp(-p3 * eta) + p4 * cp.cos(p5 * eta + p6)
        sigma0 = 10 ** (sigma0_db / 10)
        return sigma0

    @staticmethod
    def weight(wx, wy):
        return wx @ wy.T

    @staticmethod
    def vector(sx, sy):
        sy = cp.transpose(sy, (0, 1, 3, 2))
        v = sx @ sy
        return v


if __name__ == '__main__':
    AntennaParameter = {
        'Nx': 22,
        'Ny': 16,
        'Wxt': cp.ones((22, 1), dtype='complex'),
        # 'Wxt': cp.load('Wt2.npy').reshape(-1, 1),
        'Wyt': cp.ones((16, 1), dtype='complex'),
        'Wxr': cp.ones((22, 1), dtype='complex'),
        # 'Wxr': cp.load('Wt2.npy').reshape(-1, 1),
        'Wyr': cp.ones((16, 1), dtype='complex'),
        'dxRatio': 0.65,
        'dyRatio': 0.64,
    }

    RadarParameter = {
        'H': 607e3,
        'f0': 1.26e9,
        'PRF': 3250,  # test
        # 'PRF': 2000,
        'Br': 80e6,
        'N': 1    # 通道数
    }

    MissionParameter = {
        'alphaC': 26,
        'thetaSqC': 15,
        'alphaScanMin': 25.06,
        'alphaScanMax': 27.33
    }
    #
    # MissionParameter = {
    #     'alphaC': 14.3,
    #     'thetaSqC': 10,
    #     'alphaScanMin': 12.11,
    #     'alphaScanMax': 16.5
    # }

    SampleParameter = {
        'RSampleRate': 50,
        'fSampleRate': 5
    }

    Asr = ASR(AntennaParameter, RadarParameter, MissionParameter, SampleParameter)

    S0, Sa = Asr.cal_matrix()
    S0 = cp.asnumpy(S0)
    Sa = cp.asnumpy(Sa)

    # test pattern
    plt.figure()
    Asr.plot_pattern(Asr.Wxt, Asr.Wyt)
    plt.figure()
    Asr.plot_rar_r(S0, Sa, Asr.Wxr.get(), Asr.Wyt.get(), if_plot=True)
    plt.show(block=True)
    print('end')